import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

from . import tool
from matplotlib import cm, colors
from matplotlib.ticker import FormatStrFormatter
import matplotlib.patheffects as pe


def draw_def(db, cf):
    dpi = cf.DRAW_DPI
    max_px = cf.DRAW_MAX_PX
    output_png = cf.DEF_PNG

    W = max(1, db.die_ux - db.die_lx)
    H = max(1, db.die_uy - db.die_ly)
    scale = min(max_px / float(W), max_px / float(H))
    width_px = max(1, int(W * scale))
    height_px = max(1, int(H * scale))
    fig_w_in = width_px / float(dpi)
    fig_h_in = height_px / float(dpi)

    fig = plt.figure(figsize=(fig_w_in, fig_h_in), dpi=dpi)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_xlim(0, W)
    ax.set_ylim(0, H)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_axis_off()

    # 背景与边框（将外框线置于最上层并加粗）
    ax.add_patch(Rectangle((0, 0), W, H, facecolor='#F9FAFB', edgecolor='none', zorder=0))
    ax.add_patch(Rectangle((0, 0), W, H, facecolor='none', edgecolor='#FF0000', linewidth=5, zorder=100))

    # 收集需绘制的实例（用于后续重叠计算与分阶段绘制）
    cells = []
    for name, inst in db.insts.items():
        if not inst.get('placed') and not inst.get('fixed'):
            continue
        macro = inst.get('macro', '')
        w_um, h_um = db.macros_um.get(macro, (0.0, 0.0))
        w = int(round(w_um * db.dbu_per_micron))
        h = int(round(h_um * db.dbu_per_micron))
        o = inst.get('orient', 'N')
        if o in ('E', 'W', 'FE', 'FW'):
            w, h = h, w
        if w <= 0 or h <= 0:
            w = max(1, w)
            h = max(1, h)
        lx = inst.get('x', 0)
        ly = inst.get('y', 0)
        rx = lx - db.die_lx
        ry = ly - db.die_ly
        y_plot = H - (ry + h)
        # 颜色：宏单元黑色；固定的标准单元黄色；其他标准单元使用稳定映射色
        cls = db.macros_class.get(macro, '').upper()
        is_macro = (cls in ('BLOCK', 'PAD')) or (cls == '' and not macro.startswith('sky130_fd_sc_'))
        is_standard = not is_macro
        if is_macro:
            face = '#000000'
        elif is_standard and inst.get('fixed', False):
            face = '#FFFF00'
        else:
            face = tool.color_for(macro)
        cells.append({
            'name': name,
            'macro': macro,
            'rx': rx,
            'y': y_plot,
            'w': w,
            'h': h,
            'face': face
        })

    # 第一阶段：绘制填充（不画边框，避免覆盖）
    for c in cells:
        ax.add_patch(Rectangle((c['rx'], c['y']), c['w'], c['h'], facecolor=c['face'], edgecolor='none', linewidth=0.0, zorder=3))

    # 第二阶段：绘制重叠区域（红色填充），保证在填充之上、边框之下
    n = len(cells)
    for i in range(n):
        ci = cells[i]
        x1 = ci['rx']
        y1 = ci['y']
        x2 = x1 + ci['w']
        y2 = y1 + ci['h']
        for j in range(i + 1, n):
            cj = cells[j]
            xx1 = max(x1, cj['rx'])
            yy1 = max(y1, cj['y'])
            xx2 = min(x2, cj['rx'] + cj['w'])
            yy2 = min(y2, cj['y'] + cj['h'])
            if xx2 > xx1 and yy2 > yy1:
                ax.add_patch(Rectangle((xx1, yy1), xx2 - xx1, yy2 - yy1, facecolor='#FF0000', edgecolor='none', linewidth=0.0, zorder=5))

    # 第三阶段：绘制边框（置于最上层，确保不被覆盖）
    for c in cells:
        ax.add_patch(Rectangle((c['rx'], c['y']), c['w'], c['h'], facecolor='none', edgecolor='#555', linewidth=0.8, zorder=10))

    # 保存一份 PNG 供预览（项目根目录）
    fig.savefig(output_png, dpi=dpi)
    # plt.show()


def draw_congestion(db, cf):
    # 读取配置默认值
    dpi = cf.DRAW_DPI
    max_px = cf.DRAW_MAX_PX
    layer = cf.DRAW_LAYER
    alpha = cf.DRAW_ALPHA
    show_grid = cf.DRAW_SHOW_GRID
    cbar_ticks = cf.DRAW_CBAR_TICKS
    cbar_label = cf.DRAW_CBAR_LABEL
    annotate = cf.DRAW_ANNOTATE
    # 仅读取新命名：PROB_CONGESTION_TXT
    prob_path = getattr(cf, 'PROB_CONGESTION_TXT', None)

    # 解析 prob 文件
    nx, ny, utilH, utilV = parse_gcell_map(prob_path)

    # 组合层
    utilOverall = [[max(utilH[ix][iy], utilV[ix][iy]) for iy in range(ny)] for ix in range(nx)]

    # 画布设置（与 draw_def 一致坐标系）
    W = max(1, db.die_ux - db.die_lx)
    H = max(1, db.die_uy - db.die_ly)
    scale = min(max_px / float(W), max_px / float(H))
    width_px = max(1, int(W * scale))
    height_px = max(1, int(H * scale))
    fig_w_in = width_px / float(dpi)
    fig_h_in = height_px / float(dpi)

    fig = plt.figure(figsize=(fig_w_in, fig_h_in), dpi=dpi)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_xlim(0, W)
    ax.set_ylim(0, H)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_axis_off()

    # 背景与 DIE 边界
    ax.add_patch(Rectangle((0, 0), W, H, facecolor='#FFFFFF', edgecolor='none', zorder=0))
    ax.add_patch(Rectangle((0, 0), W, H, facecolor='none', edgecolor='#222', linewidth=3, zorder=100))

    # 简化绘制放置（浅色填充 + 边框）
    for _, inst in db.insts.items():
        if not inst.get('placed') and not inst.get('fixed'):
            continue
        macro = inst.get('macro', '')
        w_um, h_um = db.macros_um.get(macro, (0.0, 0.0))
        w = int(round(w_um * db.dbu_per_micron))
        h = int(round(h_um * db.dbu_per_micron))
        o = inst.get('orient', 'N')
        if o in ('E', 'W', 'FE', 'FW'):
            w, h = h, w
        w = max(1, w); h = max(1, h)
        lx = inst.get('x', 0)
        ly = inst.get('y', 0)
        rx = lx - db.die_lx
        ry = ly - db.die_ly
        y_plot = H - (ry + h)
        ax.add_patch(Rectangle((rx, y_plot), w, h, facecolor='#D0D0D0', edgecolor='#888', linewidth=0.5, zorder=2))

    # 选择要显示的利用率矩阵
    mat = utilOverall if layer == 'overall' else (utilH if layer == 'H' else utilV)
    # 颜色映射
    cmap = cm.get_cmap('RdYlGn_r')
    norm = colors.Normalize(vmin=0.0, vmax=1.2)

    # 绘制热力矩形覆盖（按 draw_def 坐标系翻转 y）
    cell_w = W / float(nx)
    cell_h = H / float(ny)
    for ix in range(nx):
        for iy in range(ny):
            u = mat[ix][iy]
            color = cmap(norm(u))
            x0 = ix * cell_w
            y0 = H - ((iy * cell_h) + cell_h)  # 翻转 y
            ax.add_patch(Rectangle((x0, y0), cell_w, cell_h, facecolor=color, edgecolor='none', alpha=alpha, zorder=50))
            # 叠加数值标签（可开关）
            if annotate:
                try:
                    ax.text(x0 + cell_w * 0.5,
                            y0 + cell_h * 0.5,
                            f"{u:.2f}",
                            ha='center', va='center',
                            fontsize=7,
                            color='#000000',
                            zorder=70,
                            path_effects=[pe.withStroke(linewidth=1.2, foreground='white')])
                except Exception:
                    pass

    # 可选：绘制 gcell 网格线
    if show_grid:
        for ix in range(1, nx):
            x = ix * cell_w
            ax.plot([x, x], [0, H], color='#999', linewidth=0.3, alpha=0.5, zorder=60)
        for iy in range(1, ny):
            y = H - iy * cell_h
            ax.plot([0, W], [y, y], color='#999', linewidth=0.3, alpha=0.5, zorder=60)

    # 颜色条（右侧）
    try:
        mappable = cm.ScalarMappable(norm=norm, cmap=cmap)
        mappable.set_array([])
        # 加宽颜色条区域，避免刻度文字被裁剪
        cax = fig.add_axes([0.965, 0.05, 0.03, 0.9])
        # 强制刻度与格式
        if cbar_ticks is None:
            cbar_ticks = [0.0, 0.25, 0.5, 0.75, 1.0, 1.2]
        cbar = fig.colorbar(mappable, cax=cax, ticks=cbar_ticks)
        # 显式设置刻度与标签，确保不同样式下也可见
        cbar.ax.set_yticks(cbar_ticks)
        cbar.ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
        cbar.ax.set_yticklabels([f"{t:.2f}" for t in cbar_ticks])
        cbar.ax.tick_params(labelsize=10, length=3, labelright=True)
        try:
            cbar.update_ticks()
        except Exception:
            pass
        if cbar_label:
            try:
                cbar.set_label(cbar_label, rotation=90, labelpad=10)
            except Exception:
                pass
    except Exception:
        pass

    # 保存输出
    out_png = cf.CONGESTION_PNG
    fig.savefig(out_png, dpi=dpi, bbox_inches='tight')
    # plt.show()


def parse_gcell_map(prob_path):
    try:
        with open(prob_path, 'r', encoding='utf-8', errors='ignore') as f:
            lines = [ln.rstrip("\n") for ln in f.readlines()]
        # 头部
        nx = ny = 0
        body_idx = 0
        for i, ln in enumerate(lines):
            if ln.startswith('#'):
                parts = ln.strip().split()
                if len(parts) >= 3 and parts[1] == 'gcell(x):':
                    try:
                        nx = int(parts[2])
                    except Exception:
                        pass
                elif len(parts) >= 3 and parts[1] == 'gcell(y):':
                    try:
                        ny = int(parts[2])
                    except Exception:
                        pass
            else:
                body_idx = i
                break
        if nx <= 0 or ny <= 0:
            print('[error] prob 头部缺少有效的 gcell 尺寸')
            return None
        # 两层数据
        layer_lines = lines[body_idx:body_idx + 2]
        if len(layer_lines) < 2:
            print('[error] prob 文件体数据不足两层')
            return None
        utilH = [[0.0 for _ in range(ny)] for _ in range(nx)]
        utilV = [[0.0 for _ in range(ny)] for _ in range(nx)]
        for z, ln in enumerate(layer_lines):
            row_tokens = ln.split('\t') if '\t' in ln else [ln]
            if len(row_tokens) != nx:
                # 兼容每行一个 ix 的写法
                flat = ln.replace('\t', ' ').split()
                print('[warn] 行分隔与预期不符，尝试兼容解析')
                # 回退：按顺序填充
                k = 0
                for ix in range(nx):
                    for iy in range(ny):
                        if k >= len(flat):
                            break
                        item = flat[k]
                        k += 1
                        if ',' in item:
                            s_str, d_str = item.split(',')[0:2]
                            try:
                                s = float(s_str)
                                d = float(d_str)
                            except Exception:
                                s = 0.0
                                d = 0.0
                            u = (d / s) if s > 0.0 else 0.0
                            if z == 0:
                                utilH[ix][iy] = u
                            else:
                                utilV[ix][iy] = u
                continue
            for ix in range(nx):
                row = row_tokens[ix].strip()
                if not row:
                    continue
                items = row.split(' ')
                for iy in range(min(ny, len(items))):
                    item = items[iy]
                    if ',' not in item:
                        continue
                    s_str, d_str = item.split(',')[0:2]
                    try:
                        s = float(s_str)
                        d = float(d_str)
                    except Exception:
                        s = 0.0
                        d = 0.0
                    u = (d / s) if s > 0.0 else 0.0
                    if z == 0:
                        utilH[ix][iy] = u
                    else:
                        utilV[ix][iy] = u
        return nx, ny, utilH, utilV
    except Exception as e:
        print(f"[error] 解析 gcell 文件失败: {e}")
        return None
